import matplotlib.pyplot as plt
import numpy as np
from get_stock import *

# 实例化
stock = get_stock()

# 策略及回测类
class strategy:
    def __init__(self):
        self.name='股票交易策略类'
    ##策略一：最简单的策略，第一次先以收盘价买入当日涨幅最大的股票，每次调仓卖出持有的全部股票，并买入近三天涨幅最大的股票
    def strategy_1(self,start_date, end_date, data, period, num):
        # 获取该时间段内所有交易日
        date_list = stock.get_trade_dates(start_date, end_date)
        # 定义一个调仓次数n
        n = len(date_list) // period
        # 对持有的股票创建一个字典
        stock_dict = {}
        for i in range(n+1):
            # 日期
            date = date_list[i * period]
            # 获取该日的股票数据df格式
            df = data[date]

            if i == 0:
                # 定义一天的涨幅diff
                df['diff'] = df['close'] - df['open']
                df = df.sort_values(by="diff", ascending=False)
                # 找到涨幅最高的股票
                df = df.head(num)
                # 获得所选股票列表
                stock_dict[date] = df.index.tolist()
            else:
                # 获得三天前的数据
                date_1 = date_list[i * period - 3]
                df_1 = data[date_1]
                df['diff'] = df['close'] - df_1['close']
                df = df.sort_values(by="diff", ascending=False)
                # 找到涨幅最高的股票
                df = df.head(num)
                # 获得每个时间点所选股票构成的字典
                stock_dict[date] = df.index.tolist()
        return stock_dict

    # 策略二：第一次先以收盘价买入当日跌幅最大的股票，之后每次调仓卖出持有的全部股票，并买入近三天跌幅最大的股票
    def strategy_2(self, start_date, end_date, data, period, num):
        # 获取该时间段内所有交易日
        date_list = stock.get_trade_dates(start_date, end_date)
        # 定义一个调仓次数n
        n = len(date_list) // period
        # 对持有的股票创建一个字典
        stock_dict = {}
        for i in range(n+1):
            # 日期
            date = date_list[i * period]
            # 获取该日的股票数据df格式
            df = data[date]

            if i == 0:
                # 定义一天的涨幅diff
                df['diff'] = df['close'] - df['open']
                df = df.sort_values(by="diff", ascending=False)
                df=df.fillna(1)
                # 找到跌幅最大的股票
                df = df.tail(num)
                # 获得所选股票列表
                stock_dict[date] = df.index.tolist()
            else:
                # 获得三天前的数据
                date_1 = date_list[i * period - 3]
                df_1 = data[date_1]
                df['diff'] = df['close'] - df_1['close']
                df = df.sort_values(by="diff", ascending=False)
                df=df.fillna(1)
                # 找到跌幅最大的股票
                df = df.tail(num)
                # 获得每个时间点所选股票构成的字典
                stock_dict[date] = df.index.tolist()
        return stock_dict

   # 使用移动均线交叉策略对多只股票量化分析
    def strategy_3(self, start_date, end_date, code_list):
        # 定义一个字典用来存储code及对应的交易时间
        dict={}
        for code in code_list:
            data = stock.get_a_stock(code, start_date, end_date)
            trade_time = self.mean_reversion(data)
            dict[code]=trade_time
        return dict

    # 使用移动均线交叉策略对单只股票量化分析
    def mean_reversion(self, data):
        # 初始化短期和长期窗口
        short_window = 40
        long_window = 100
        # 初始化 `signals` 数据框，增加 `signal` 列
        signals = pd.DataFrame(index=data.index)
        signals['signal'] = 0.0
        # 创建短期简单移动均值
        signals['short_mavg'] = data['close'].rolling(window=short_window, min_periods=1, center=False).mean()
        # 创建长期简单移动均值
        signals['long_mavg'] = data['close'].rolling(window=long_window, min_periods=1, center=False).mean()
        # 生成信号
        signals['signal'][short_window:] = np.where(signals['short_mavg'][short_window:]
                                                    > signals['long_mavg'][short_window:], 1.0, 0.0)
        # 生成交易命令
        signals['positions'] = signals['signal'].diff()
        buy_time = signals.loc[signals.positions == 1.0].index.tolist()
        sell_time = signals.loc[signals.positions == -1.0].index.tolist()
        # print('买入时间：')
        # print(buy_time)  # 字符串列表类型                                                                                               '').split(',')
        # print('卖出时间：')
        # print(sell_time)
        trade_time = [buy_time, sell_time]

        # 可视化
        # 初始化图形
        fig = plt.figure(figsize=(12, 8))
        # 增加子图并设置y轴标签
        ax1 = fig.add_subplot(111, ylabel='Price')
        # 将object类型转换为numeric以便绘图
        data['close'] = (data['close']).astype(float)
        signals['short_mavg'] = (signals['short_mavg']).astype(float)
        signals['long_mavg'] = (signals['long_mavg']).astype(float)
        # 绘制收盘价曲线
        data['close'].plot(ax=ax1, color='r', lw=2.)
        # 绘制短期和长期移动均线
        signals[['short_mavg', 'long_mavg']].plot(ax=ax1, lw=2.)
        # 绘制买入信号  标记点位置有点问题
        # b = data['close'][signals.loc[signals.positions == 1.0].index]
        # ax1.plot(b,'^', markersize=10, color='m')
        # 绘制卖出信号  标记点位置有点问题
        # ax1.plot(signals.loc[signals.positions == -1.0].index,
        #          signals.short_mavg[signals.positions == -1.0],
        #          'v', markersize=10, color='k')
        # 显示做图
        plt.show()
        return trade_time

    ## 使用移动均线交叉策略对单只股票量化回测,返回交易的时间
    def backtest(self, code, start_date, end_date):
        data = stock.get_a_stock(code, start_date, end_date)
        trade_time=self.mean_reversion(data)
        print('买入时间：')
        print(trade_time[0])
        print('卖出时间：')
        print(trade_time[1])





